package com.jianjang.llm.core.aspect;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.servlet.http.HttpServletRequest;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.context.request.async.DeferredResult;
import org.springframework.web.context.request.async.WebAsyncTask;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;

@Aspect
@Component
public class ControllerLoggingAspect {

    private static final Logger log = LoggerFactory.getLogger(ControllerLoggingAspect.class);

    private final ObjectMapper objectMapper;

    public ControllerLoggingAspect(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    @Around("within(com.jianjang.llm.controller..*)")
    public Object logController(ProceedingJoinPoint joinPoint) throws Throwable {
        long startAtMs = System.currentTimeMillis();
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        String methodPath = signature.getDeclaringTypeName() + "." + signature.getName();

        HttpServletRequest request = currentRequest();
        String httpMethod = request != null ? request.getMethod() : "N/A";
        String uri = request != null ? request.getRequestURI() : "N/A";
        String queryString = request != null ? request.getQueryString() : null;

        Map<String, Object> requestParams = request != null ? extractRequestParams(request) : new HashMap<>();
        String requestBodyJson = extractRequestBodyJson(joinPoint, request);

        log.info("[REQ] {} {}{} method={} params={} body={}",
                httpMethod,
                uri,
                queryString != null ? ("?" + queryString) : "",
                methodPath,
                safeToJson(requestParams),
                requestBodyJson);

        CountDownLatch latch = new CountDownLatch(1);
        try {
            Object result = joinPoint.proceed();

            if (result instanceof DeferredResult<?> deferred) {
                deferred.onCompletion(() -> {
                    long costMs = System.currentTimeMillis() - startAtMs;
                    log.info("[RESP] {} {} method={} cost={}ms data={}", httpMethod, uri, methodPath, costMs, safeToJson(deferred.getResult()));
                    latch.countDown();
                });
                return result;
            }

            if (result instanceof WebAsyncTask<?> webTask) {
                webTask.onCompletion(() -> {
                    long costMs = System.currentTimeMillis() - startAtMs;
                    log.info("[RESP] {} {} method={} cost={}ms", httpMethod, uri, methodPath, costMs);
                    latch.countDown();
                });
                return webTask;
            }

            if (result instanceof Callable<?> callable) {
                Callable<?> wrapped = () -> {
                    try {
                        Object value = callable.call();
                        long costMs = System.currentTimeMillis() - startAtMs;
                        log.info("[RESP] {} {} method={} cost={}ms data={}", httpMethod, uri, methodPath, costMs, safeToJson(value));
                        return value;
                    } catch (Throwable t) {
                        long costMs = System.currentTimeMillis() - startAtMs;
                        log.error("[RESP] {} {} method={} cost={}ms error={}", httpMethod, uri, methodPath, costMs, t.getMessage(), t);
                        throw t;
                    } finally {
                        latch.countDown();
                    }
                };
                return wrapped;
            }

            long costMs = System.currentTimeMillis() - startAtMs;
            log.info("[RESP] {} {} method={} cost={}ms data={}", httpMethod, uri, methodPath, costMs, safeToJson(result));
            latch.countDown();
            return result;
        } catch (Throwable t) {
            long costMs = System.currentTimeMillis() - startAtMs;
            log.error("[RESP] {} {} method={} cost={}ms error={}", httpMethod, uri, methodPath, costMs, t.getMessage(), t);
            latch.countDown();
            throw t;
        }
    }

    private HttpServletRequest currentRequest() {
        RequestAttributes attrs = RequestContextHolder.getRequestAttributes();
        if (attrs instanceof ServletRequestAttributes sra) {
            return sra.getRequest();
        }
        return null;
    }

    private Map<String, Object> extractRequestParams(HttpServletRequest request) {
        Map<String, String[]> raw = request.getParameterMap();
        Map<String, Object> flat = new HashMap<>();
        for (Map.Entry<String, String[]> e : raw.entrySet()) {
            String key = e.getKey();
            String[] values = e.getValue();
            if (values == null) {
                flat.put(key, null);
            } else if (values.length == 1) {
                flat.put(key, values[0]);
            } else {
                flat.put(key, Arrays.asList(values));
            }
        }
        return flat;
    }

    private String extractRequestBodyJson(ProceedingJoinPoint joinPoint, HttpServletRequest request) {
        try {
            // Prefer @RequestBody parameter if present (covers application/json)
            MethodSignature signature = (MethodSignature) joinPoint.getSignature();
            Method method = signature.getMethod();
            Annotation[][] paramAnns = method.getParameterAnnotations();
            Object[] args = joinPoint.getArgs();
            for (int i = 0; i < paramAnns.length && i < args.length; i++) {
                boolean hasRequestBody = Arrays.stream(paramAnns[i])
                        .anyMatch(a -> a.annotationType() == RequestBody.class);
                if (hasRequestBody) {
                    return safeToJson(args[i]);
                }
            }

            // Fallback: if content-type is JSON and no @RequestBody found, serialize non-framework args
            if (request != null && request.getContentType() != null
                    && request.getContentType().startsWith(MediaType.APPLICATION_JSON_VALUE)) {
                Object[] allArgs = joinPoint.getArgs();
                Object candidate = Arrays.stream(allArgs)
                        .filter(o -> o != null && !isFrameworkType(o))
                        .findFirst()
                        .orElse(null);
                return safeToJson(candidate);
            }
        } catch (Exception ignore) {
        }
        return null;
    }

    private boolean isFrameworkType(Object o) {
        String pkg = o.getClass().getPackageName();
        return pkg.startsWith("org.springframework") || pkg.startsWith("jakarta.") || pkg.startsWith("javax.");
    }

    private String safeToJson(Object obj) {
        if (obj == null) {
            return null;
        }
        try {
            String json = objectMapper.writeValueAsString(obj);
            return truncate(json, 4000);
        } catch (JsonProcessingException e) {
            // Fallback to toString
            return truncate(String.valueOf(obj), 4000);
        }
    }

    private String truncate(String s, int maxLen) {
        if (s == null || s.length() <= maxLen) {
            return s;
        }
        return s.substring(0, maxLen) + "…(truncated)";
    }
}


